/**
* Copyright (c) [2023] [Tenny]
   [mysql-migrate] is licensed under Mulan PSL v2.
   You can use this software according to the terms and conditions of the Mulan PSL v2.
   You may obtain a copy of Mulan PSL v2 at:
            http://license.coscl.org.cn/MulanPSL2
   THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
   See the Mulan PSL v2 for more details.
*/
use chrono::format::strftime::StrftimeItems;
use chrono::Local;
use clap::{arg, ArgMatches, Command};
use mysql::prelude::Queryable;
use mysql::{Pool, PooledConn, TxOpts};
use regex::Regex;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::env;
use std::error::Error;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::string::String;

fn cli() -> Command {
    Command::new("migrate")
        .version("0.0.1")
        .author("Tenny. <tenny.shu@foxmail.com>")
        .about("mysql数据库迁移工具cli")
        .subcommand_required(true)
        .arg_required_else_help(true)
        .subcommands([
            Command::new("tables")
                .about("显示所有的表")
                .arg(arg!(-e --env [env_names] "环境变量文件列表")),
            Command::new("columns")
                .about("显示表结构")
                .arg_required_else_help(true)
                .arg(arg!(<name> "表名"))
                .arg(arg!(-e --env [env_names] "环境变量文件列表")),
            Command::new("run")
                .about("执行数据库迁移")
                .arg(arg!(-e --env [env_names] "环境变量文件列表")),
        ])
}

/// 从指定目录的环境变量文件中加载环境变量并返回一个包含键值对的 HashMap。
/// 如果未提供目录，则默认为当前命令运行的目录。
/// 如果未提供环境变量文件列表，则将使用默认的文件名列表: `[".env", ".env.local", ".env.development", ".env.production"]`
/// 并且环境变量文件是有序的, 后面的文件的内容会替换前面的文件的内容
///
/// # 参数
///
/// - `dir`: 环境变量文件目录的可选路径。
/// - `files`: 环境变量文件名称列表的可选数组。
///
/// # 返回值
///
/// 返回一个包含加载的环境变量的 HashMap，其中键和值都以字符串形式表示。
///
/// # 示例
///
/// ```
/// use std::collections::HashMap;
///
/// let env_dir = Some("/path/to/env/files".to_string());
/// let env_files = Some([".env", ".env.local"]);
/// let env_map: HashMap<String, String> = load_env(env_dir, env_files);
/// ```
/// 也可以不使用返回值, 直接通过在 `env` 中获取
/// ```
/// let a = env::var("key").unwrap_or_else(|_| String::from("s"));
/// ```
fn load_env(dir: Option<String>, files: Option<Vec<&str>>) -> HashMap<String, String> {
    let mut envs = HashMap::new();
    // 环境变量文件目录; 如果没有传环境变量文件目录, 则默认为当前命令运行目录
    let env_dir = dir.unwrap_or_else(|| env::current_dir().unwrap().display().to_string());
    let env_files =
        files.unwrap_or_else(|| vec![".env", ".env.local", ".env.development", ".env.production"]);
    for &envfile in &env_files {
        // 文件完整路径
        let mut envpath = PathBuf::new();
        envpath.push(&env_dir);
        envpath.push(&envfile);
        // 遍历读取文件每一行
        if let Ok(file) = File::open(&envpath) {
            let f = BufReader::new(file);
            for line in f.lines() {
                if let Ok(line) = line {
                    // 将每一行的数据按照 - 拆分为两段
                    let parts: Vec<&str> = line.splitn(2, "=").collect();
                    if parts.len() == 2 {
                        let key = parts[0].trim().to_string();
                        // 不是注释行
                        if key.starts_with("#") == false {
                            let value = parts[1].trim().to_string();
                            envs.insert(key.clone(), value.clone());

                            env::set_var(key, value)
                        }
                    }
                }
            }
        }
    }
    return envs;
}

/// 记录日志并打印到标准输出
///
/// # 参数
///
/// - `msg`: 要记录的日志消息
/// - `level`: 日志级别，0 表示调试，2 表示警告，3 表示错误，其他为信息
///
/// # 示例
///
/// ```rust
/// log_record("这是一条日志消息", 2);
/// ```
fn log_record(msg: &str, level: i8) {
    // 匹配日志级别，确定级别字符串和颜色代码
    let (level_str, color_code) = match level {
        0 => ("debug", 4),
        2 => ("warning", 3),
        3 => ("error", 1),
        _ => ("info", 2),
    };
    // 获取当前时间并格式化为字符串
    let dt_str = Local::now()
        .format_with_items(StrftimeItems::new("%Y-%m-%d %H:%M:%S%.3f"))
        .to_string();
    // 打印带有颜色的日志信息
    println!(
        "\x1B[3{}m[{}] [{}] - {}\x1B[0m",
        color_code,
        dt_str,
        level_str.to_uppercase(),
        msg
    );
}

/// 加载配置
fn load_config(sub_matches: &ArgMatches) -> HashMap<String, String> {
    // 传递的环境变量文件列表
    let env_names = sub_matches.get_one::<String>("env").map(|s| s.as_str());
    let mut env_vec: Option<Vec<&str>> = None;
    // 如果传递的环境变量文件列表不为空，则将其转换为字符串数组
    if env_names.is_some() {
        env_vec = Some(env_names.unwrap().split(",").collect::<Vec<&str>>());
    }
    // 返回加载环境变量的结果
    return load_env(None, env_vec);
}

/// 查询并记录数据库中的表名列表。
///
/// 该函数从给定的数据库连接池中获取连接，执行 SQL 查询以获取表名列表，
/// 然后将格式化后的表名列表记录到日志中。
///
/// # 参数
///
/// - `pool`: 数据库连接池，用于获取数据库连接
///
/// # 示例
///
/// ```rust
/// let pool = create_database_pool();
/// show_tables(pool);
/// ```
fn show_tables(pool: Pool) {
    let mut conn = pool.get_conn().unwrap();
    let res: Vec<String> = conn.query("SHOW TABLES").unwrap();
    let mut tables: Vec<String> = res.iter().map(|item| format!("- {0}", item)).collect();
    tables.insert(0, String::from("TABLES: "));
    log_record(&tables.join("\r\n"), 1);
}

/// 从数据库获取并显示表的列信息。
///
/// 此函数使用提供的连接池连接到数据库，检索指定表的列信息。然后构造并记录列信息的格式化表示。
///
/// # 参数
///
/// * `pool` - 一个 MySQL 数据库连接池。
/// * `name` - 要获取列信息的表名。
///
fn show_table_columns(pool: Pool, name: String) {
    let mut conn = pool.get_conn().unwrap();

    let mut infos: Vec<String> = vec![
        format!("table `{}` columns: ", name),
        "字段 -- 类型 -- 空 -- 索引 -- 默认 -- 其它".to_string(),
    ];
    conn.query_iter(format!("SHOW COLUMNS from `{}`", name))
        .unwrap()
        .for_each(|row| {
            let r = row.unwrap();
            let mut column_values: Vec<String> = vec![];
            // 遍历每一列并获取值
            for i in 0..r.columns().len() {
                let column_value: String =
                    r.get_opt(i).unwrap().unwrap_or_else(|_| String::from(""));
                column_values.push(column_value);
            }
            infos.push(column_values.join(" -- "));
        });
    log_record(infos.join("\r\n").as_str(), 1);
}

/// 创建数据库连接池并返回结果, 并且会打印相关连接信息。
///
/// # 参数
///
/// * `url` - 数据库连接 URL
///
/// # 返回值
///
///如果连接池创建成功，返回一个 `Result`，其中包含数据库连接池，否则返回一个包含错误信息的 `Result`。
///
/// # 示例
///
/// ```
/// let pool = create_pool("").unwrap();
/// let conn = pool.get_conn();
/// ```
fn create_pool(url: &str) -> Result<Pool, Box<dyn Error>> {
    match Pool::new(url) {
        Ok(pool) => {
            log_record("连接成功", 1);
            return Ok(pool);
        }
        Err(e) => {
            log_record(&e.to_string(), 3);
            Err(Box::new(e))
        }
    }
}

/// 解析 SQL 文件并将其中的 SQL 查询语句提取出来。
///
/// # 参数
///
/// * `sql_path` - 包含 SQL 查询语句的文件的路径。
///
/// # 返回值
///
/// 返回一个包含所有提取出的 SQL 查询语句的字符串向量。
fn parse_sql_file(sql_path: PathBuf) -> Vec<String> {
    let mut all_sqls: Vec<String> = vec![];
    // 通过正则移除注释
    let pattern = Regex::new(r"/\*.*?\*/").unwrap();
    if let Ok(file) = File::open(sql_path) {
        let reader = BufReader::new(file);
        let mut item_sql: Vec<String> = vec![];
        for line in reader.lines() {
            if let Ok(line_content) = line {
                let sql_line_content = pattern.replace_all(&line_content, "").trim().to_string();
                if sql_line_content == ";" {
                    continue;
                }
                item_sql.push(sql_line_content.to_string());
                if sql_line_content.ends_with(";") {
                    all_sqls.push(item_sql.join(""));
                    item_sql.clear();
                }
            }
        }
    }
    return all_sqls;
}

/// SQL 文件信息
struct SqlFileInfoStruct {
    /// 版本
    version: String,
    /// 描述
    description: String,
    /// 文件名称
    file_name: String,
    /// 文件内容
    content: Vec<String>,
}

/// 批量插入迁移记录数据列表。
///
/// 这个函数接受一个数据库连接（`conn`）、目标表名（`table`）、列名列表（`columns`）和一个包含要插入的数据的元组向量（`values`）作为参数。
///
/// # 参数
///
/// - `conn`: 数据库连接对象，通常是从数据库连接池中获取的。
/// - `table`: 目标数据表的名称。
/// - `columns`: 列名列表，包含要插入的数据的列。
/// - `values`: 包含要插入的数据的元组向量，每个元组包括多个字段的值。
///
/// # 注意事项
///
/// - 这个函数将生成一个SQL批量插入语句并将其执行在数据库上。
/// - 函数执行成功时不返回任何值，如果执行失败会产生错误。
///
/// # 示例
///
/// ```rust
/// use mysql::PooledConn;
///
/// let conn = // 获取数据库连接的代码;
/// let table = "my_table".to_string();
/// let columns = vec!["col1".to_string(), "col2".to_string(), "col3".to_string()];
/// let values = vec![
///     ("value1".to_string(), "value2".to_string(), "value3".to_string(), 42, 1000, 1),
///     // 更多数据行...
/// ];
///
/// batch_insert_schema(conn, table, columns, values);
/// ```
fn batch_insert_schema(
    mut conn: PooledConn,
    table: String,
    columns: Vec<String>,
    values: Vec<(String, String, String, String, i64, i8)>,
) {
    let mut insert_schema_sqls = vec![format!("insert into `{}`", table)];
    let mut insert_column_schemas: Vec<String> = vec![];
    for column in columns {
        insert_column_schemas.push(format!("`{}`", column));
    }
    insert_schema_sqls.push(format!("({})", insert_column_schemas.join(",")));
    insert_schema_sqls.push("values".to_string());
    let mut value_schemas: Vec<String> = vec![];
    for value in values {
        value_schemas.push(format!(
            "(\"{}\", \"{}\", \"{}\", \"{}\", {}, {})",
            value.0, value.1, value.2, value.3, value.4, value.5
        ))
    }
    insert_schema_sqls.push(value_schemas.join(","));
    let batch_sql = insert_schema_sqls.join(" ");
    log_record(batch_sql.as_str(), 1);
    conn.exec_drop(batch_sql, ()).unwrap();
}

/// 批量修改迁移历史记录
fn batch_update_schema(
    mut conn: PooledConn,
    table: String,
    values: Vec<(String, String, String, String, i64, i8)>,
) {
    let mut updates: Vec<String> = vec![];
    for value in values {
        let update_sql = format!(
            "update `{}` set `description`=\"{}\",`script`=\"{}\",`run_time`=\"{}\",`execution_time`={},`success`={} where `version`=\"{}\"",
            table, value.1, value.2, value.3, value.4, value.5, value.0
        );
        conn.exec_drop(update_sql.clone(), ()).unwrap();
        updates.push(update_sql);
    }
    log_record(updates.join(",").as_str(), 1);
}

fn run_migrate(
    pool: Pool,
    migrate_table_name: String,
    migrate_file_dir: String,
    last_line: String,
) -> String {
    let mut conn = pool.get_conn().unwrap();

    // 创建迁移记录表
    conn.query_drop(format!(
        r#"
    CREATE TABLE IF NOT EXISTS `{}` (
        `id` INT NOT NULL AUTO_INCREMENT,
        `version` VARCHAR(50) NULL DEFAULT NULL COLLATE 'utf8mb4_general_ci' COMMENT '版本号',
        `description` VARCHAR(200) NULL DEFAULT NULL COLLATE 'utf8mb4_general_ci' COMMENT '描述',
        `script` VARCHAR(1000) NOT NULL COLLATE 'utf8mb4_general_ci' COMMENT '文件名称',
        `run_time` TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP COMMENT '命令运行时间',
        `execution_time` INT NOT NULL DEFAULT 1 COMMENT '执行花费时间(单位: 毫秒)',
        `success` TINYINT NOT NULL DEFAULT 1 COMMENT '状态,1-执行成功,0-执行失败',
        PRIMARY KEY (`id`) USING BTREE,
        UNIQUE INDEX `version_idx` (`version`) USING BTREE,
        INDEX `success_idx` (`success`) USING BTREE
    )
    COLLATE='utf8mb4_general_ci'
    COMMENT='数据库迁移版本管理表'
    ENGINE=InnoDB;
    "#,
        migrate_table_name
    ))
    .unwrap();

    // 校验迁移文件目录是否存在
    let mut migration_dir = PathBuf::from(env::current_dir().unwrap());
    migration_dir.push(Path::new(migrate_file_dir.as_str()));

    let mut update_last_line = last_line.clone(); // 修改后的最新版本号

    // 如果迁移文件目录存在才进行迁移
    if migration_dir.exists() {
        let files_entry = migration_dir.read_dir().unwrap();
        let mut contents: Vec<SqlFileInfoStruct> = vec![];
        let mut versions: Vec<String> = vec![];
        for file_entry_res in files_entry {
            let filepath = file_entry_res.unwrap().path();
            if filepath.is_file() {
                let mut file_info = SqlFileInfoStruct {
                    version: "".to_string(),
                    file_name: "".to_string(),
                    description: "".to_string(),
                    content: vec![],
                };
                // 解析文件内容
                let file_sqls = parse_sql_file(filepath.clone());
                // 文件名带后缀: v20231231184101__Create_Table.sql
                let file_name = filepath.file_name().unwrap().to_str().unwrap();
                // 文件名不带后缀: v20231231184101__Create_Table
                let file_stem = filepath.file_stem().unwrap().to_str().unwrap();
                let infos: Vec<&str> = file_stem.splitn(2, "__").collect();
                // 版本信息: 20231231184101
                let mut version = infos[0].to_string();
                if version.starts_with("v") || version.starts_with("V") {
                    version = version[1..].to_string()
                }
                // 只有超过上一次成功的最大的版本, 才迁移
                match version.partial_cmp(&last_line) {
                    Some(Ordering::Greater) => {
                        // 描述信息: Create_Table
                        let mut description = "";
                        if infos.len() == 2 {
                            description = infos[1];
                        }

                        versions.push(version.clone());

                        file_info.version = version;
                        file_info.description = description.to_string();
                        file_info.file_name = file_name.to_string();
                        file_info.content = file_sqls;
                        contents.push(file_info);
                    }
                    None | Some(_) => {}
                }
            }
        }
        if contents.len() > 0 {
            // 按版本号进行排序
            contents.sort_by(|a, b| a.version.partial_cmp(&b.version).unwrap());
            // 查询将要执行且已经执行过的版本迁移
            let sql = format!(
                "select `version`, `success` from `{}` where `version` in ({})",
                migrate_table_name,
                versions.join(", ")
            );
            let runed_list: Vec<(String, i8)> = conn
                .query_map(sql, |(version, success)| (version, success))
                .unwrap();
            // 已执行且成功, 已执行且失败
            let mut runed_success: Vec<String> = vec![];
            let mut runed_fail: Vec<String> = vec![];
            for run_item in runed_list.iter() {
                if run_item.1 == 1 {
                    // 执行成功
                    runed_success.push(run_item.0.to_string());
                } else {
                    runed_fail.push(run_item.0.to_string());
                }
            }
            // 执行后新增和修改的 schema 表的记录
            let mut insert_schema: Vec<(String, String, String, String, i64, i8)> = vec![];
            let mut update_schema: Vec<(String, String, String, String, i64, i8)> = vec![];
            // 是否是连续的最新版本, 连续则最新版本覆盖, 否则不做任何修改
            let mut is_continue_line = true;
            // 遍历每个文件的执行命令
            for runing_item in contents {
                let version = runing_item.version.clone();
                // 执行成功的列表里面不包含当前版本, 则执行迁移
                if !runed_success.contains(&version) {
                    let run_start_time = Local::now()
                        .format_with_items(StrftimeItems::new("%Y-%m-%d %H:%M:%S%.3f"))
                        .to_string();
                    let start_time = Local::now().timestamp_millis();
                    // 每一个文件内容使用事务执行
                    let mut tx = conn.start_transaction(TxOpts::default()).unwrap();
                    let mut is_success = 1; // 迁移是否成功, 1 - 成功, 0 - 失败
                    for script_item in runing_item.content {
                        is_success = match tx.exec_drop(script_item, ()) {
                            Ok(()) => 1,
                            Err(err) => {
                                log_record(
                                    format!(
                                        "{} run faied: {}",
                                        runing_item.file_name,
                                        &err.to_string()
                                    )
                                    .as_str(),
                                    3,
                                );
                                0
                            }
                        };
                        if is_success == 0 {
                            break;
                        }
                    }
                    if is_success == 1 {
                        tx.commit().unwrap_or_default();
                        log_record(
                            format!("{} run success!", runing_item.file_name).as_str(),
                            1,
                        );
                        if is_continue_line {
                            update_last_line = version.clone();
                        }
                    } else {
                        tx.rollback().unwrap_or_default();
                        is_continue_line = false;
                    }
                    let end_time = Local::now().timestamp_millis();
                    let time_diff = end_time - start_time;
                    let runed_item = (
                        version,
                        runing_item.description,
                        runing_item.file_name,
                        run_start_time,
                        time_diff,
                        is_success,
                    );
                    if runed_fail.contains(&runing_item.version) {
                        update_schema.push(runed_item);
                    } else {
                        insert_schema.push(runed_item);
                    }
                }
            }
            // 插入新增的迁移记录
            if insert_schema.len() > 0 {
                batch_insert_schema(
                    conn,
                    migrate_table_name.clone(),
                    vec![
                        "version".to_string(),
                        "description".to_string(),
                        "script".to_string(),
                        "run_time".to_string(),
                        "execution_time".to_string(),
                        "success".to_string(),
                    ],
                    insert_schema,
                );
            }
            // 修改之前错误现在执行正确的记录
            if update_schema.len() > 0 {
                batch_update_schema(pool.get_conn().unwrap(), migrate_table_name, update_schema);
            }
        }
    }
    return update_last_line;
}

/// 读取缓存最新执行记录版本, 并保存到 HashMap 中
///
/// 示例:
///
/// ```rust
/// let last_map = read_last_lines();
/// ```
fn read_last_lines() -> HashMap<String, String> {
    let mut line_tmp_path = PathBuf::from(env::temp_dir());
    line_tmp_path.push("migration_last_lines.txt");
    let mut last_map: HashMap<String, String> = HashMap::new();
    if let Ok(file) = File::open(line_tmp_path) {
        let reader = BufReader::new(file);
        for line in reader.lines() {
            if let Ok(line_content) = line {
                let contents: Vec<&str> = line_content.splitn(2, "=").collect();
                if contents.len() == 2 {
                    last_map.insert(contents[0].to_string(), contents[1].to_string());
                }
            }
        }
    }
    return last_map;
}

/// 将最新迁移版本信息写入缓存
///
/// 示例:
///
/// ```rust
/// write_last_lines(map);
/// ```
fn write_last_lines(line_map: HashMap<String, String>) {
    let mut line_tmp_path = PathBuf::from(env::temp_dir());
    line_tmp_path.push("migration_last_lines.txt");
    if let Ok(file) = File::create(line_tmp_path) {
        let mut writer: BufWriter<File> = BufWriter::new(file);
        for (key, value) in line_map {
            let content_bytes = format!("{}={}", key, value);
            writer.write_all(content_bytes.as_bytes()).unwrap();
            writer.write_all(b"\n").unwrap(); // 写入换行符
        }
        writer.flush().unwrap();
    }
}

fn main() {
    let matches = cli().get_matches();

    match matches.subcommand() {
        // 显示所有的表
        Some(("tables", sub_matches)) => {
            // 加载配置
            let configs = load_config(sub_matches);
            // 数据库连接池
            let pool = create_pool(configs.get("MIG_MYSQL_URL").unwrap().as_str()).unwrap();
            // 显示所有的表名
            show_tables(pool);
        }
        // 显示表结构
        Some(("columns", sub_matches)) => {
            // 加载配置
            let configs = load_config(sub_matches);
            // 数据库连接池
            let pool = create_pool(configs.get("MIG_MYSQL_URL").unwrap().as_str()).unwrap();
            // 显示表结构
            show_table_columns(
                pool,
                sub_matches.get_one::<String>("name").unwrap().to_string(),
            );
        }
        Some(("run", sub_matches)) => {
            // 加载配置
            let configs = load_config(sub_matches);
            // 最新加载信息
            let mut last_lines = read_last_lines();
            // 数据库连接池
            let pool = create_pool(configs.get("MIG_MYSQL_URL").unwrap().as_str()).unwrap();
            // 迁移目录
            let mig_dir = match configs.get("MIG_LOCATION") {
                Some(v) => v.to_string(),
                None => "migrations".to_string(),
            };
            // 是否启用迁移
            let enable = match configs.get("MIG_ENABLED") {
                Some(v) => v.to_string(),
                None => "1".to_string(),
            };
            // 迁移历史的记录表格名称
            let schema = match configs.get("MIG_NAME") {
                Some(v) => v.to_string(),
                None => "mm".to_string(),
            };
            // 最新一次成功的迁移版本
            let mut last_line = match last_lines.get(&schema) {
                Some(last_line_version) => last_line_version.to_string(),
                None => "1".to_string(),
            };
            let mig_table_name = format!("{}_migrate_schema_history", schema);
            if enable == "1" {
                // 执行迁移
                last_line = run_migrate(pool, mig_table_name, mig_dir, last_line);
            }
            // 写入最新的行
            last_lines.insert(schema, last_line);
            write_last_lines(last_lines);
        }
        Some(_) | None => todo!(),
    }
}
